## 
rm(list = ls())

library(tidyverse)
library(lfe)
library(pbapply)
library(broom)

## Get the data

census <- read_rds('data/muni_finance.rds')%>% 
  mutate(state_id = factor(state_id)) %>% 
  mutate(east = ifelse(as.numeric(substr(ags, 1, 2)) > 11, 1, 0))

## Rename states
levels(census$state_id) <- c('S-H', 'HH', 'NDS', 'Bremen', 'NRW',
                             'Hessen', 'RP', 'BW', 'Bayern', 'Saarland',
                             'Berlin', 'Brandenburg', 'M-V', 'Sachsen',
                             'Sachsen-Anhalt', 'Thueringen')

## Normalize variables by the pre-census population

bt_tvars <- str_detect(colnames(census), '_total_') %>%
  colnames(census)[.] 

## Do

census[, bt_tvars] <- sapply(bt_tvars, function(v) {
  census[, v] / census$bt_pop_2012
})

## To long

plot_df <- census %>% 
  pivot_longer(cols = matches('bt_'), names_sep = '_(?!.*_)',
               names_to = c('name', 'year')) %>% 
  pivot_wider(names_from = 'name', values_from = 'value') %>% 
  mutate(year = as.numeric(year)) %>%
  filter(between(pop_dec_09, 0, 20000))

## Create treatment and post indicators

regdata <- plot_df %>%
  mutate(post = ifelse(year > 2013, 1, 0),
         treated_post = (as.numeric(treated) - 1) * post)

## Get info on census application and timing

states <- read_rds("data/states_census.rds") %>%
  dplyr::select(state_name, applies_census, census_first_year) %>%
  mutate(state_name = dplyr::recode(state_name, `Schleswig-Holstein` = 'S-H',
                                    `Niedersachsen` = 'NDS',
                                    `Nordrhein-Westfalen` = 'NRW',
                                    `Thüringen` = 'Thueringen',
                                    `Baden-Württemberg` = 'BW',
                                    `Mecklenburg-Vorpommern` = 'M-V',
                                    `Rheinland-Pfalz` = 'RP'))

## Merge to DiD file

regdata <- left_join(regdata, states, 
                     by = c('state_id' = 'state_name')) 

## Time relative to treatment

regdata <- regdata %>%
  filter(!is.na(year)) %>% 
  mutate(time_rel = year - census_first_year,
         time_rel = ifelse(time_rel > -1, time_rel + 1, time_rel),
         post_rel = ifelse(time_rel > -1, 1, 0),
         treated_post_rel = post_rel * (as.numeric(treated) - 1))

## Declare outcomes

outcomes <- c("bt_spending_total", "bt_ssz_total", "bt_current_debt_total")

## Gen state level trends

regdata <- regdata %>% 
  mutate(time_rel_state = paste0(time_rel, state_id),
         time_rel_county = paste0(time_rel, county))

#### Conditional on time periods pre/post census ####

periods = 1:3
bw_vec <- seq(500, 2500, 500)
exclude_bw = c(T,F)
subset_list = list(regdata, regdata %>% filter(east == 1),
                   regdata %>% filter(east == 0))
subset_labels = c("Full", "East", "West")

out_sens_p <- pblapply(bw_vec, function(i) {
  lapply(periods, function(p) {
    lapply(outcomes, function(o) {
      lapply(1:length(subset_list), function(j) {
        
        ## Get subset
        
        data_sub <- subset_list[[j]]
        label_sub <- subset_labels[j]
        
        data_sub[, 'outcome'] <- data_sub[, o]
        
        ## Run
        
        m4 <- felm(outcome ~ treated_post_rel | ags + time_rel | 0 | ags,
                   data = data_sub %>% filter(between(pop_dec_09, 10000 - i,
                                                      10000 + i)) %>%
                     filter(!state_id == 'B-W'), 
                   subset = applies_census == 1 & 
                     between(time_rel, -p,
                             p))
        
        ## Get coef
        out <- broom::tidy(m4, conf.int = T) %>%
          mutate(bw = i, period = p, outcome = o,
                 sample = label_sub)
        out
      }) %>% reduce(rbind)
    }) %>% reduce(rbind)
  }) %>% reduce(rbind)
}) %>% reduce(rbind) %>% 
  mutate(outcome = str_remove(outcome, 'bt_|_total'))

## 90% CIs

out_sens_p <- out_sens_p %>% 
  mutate(conf.low90 = estimate  - qnorm(0.95) * std.error,
         conf.high90 = estimate  + qnorm(0.95) * std.error)

## Select BWs

bw_vec <- seq(500, 2500, 500)
pd <- position_dodge(0.4)

## Plot for paper 

plot_paper <- out_sens_p %>% 
  filter(period < 4) %>% 
  filter(outcome == 'spending_total') %>% 
  mutate(period = paste0('Years included pre/post treatment: ', period)) %>% 
  filter(sample == "Full")

# Figure 3 - Effect on spending ----

p1 <- ggplot(plot_paper %>% filter(str_detect(period, '1')), 
             aes(bw, estimate)) +
  geom_hline(yintercept = 0, linetype = 'dotted') +
  geom_errorbar(aes(ymin = conf.low, ymax = conf.high),
                width = 0, size = 0.3, position = pd) +
  geom_errorbar(aes(ymin = conf.low90,
                    ymax = conf.high90),
                width = 0, size = 0.8,
                position = pd) +
  geom_point(shape = 21, fill = 'white', size = 2,
             position = pd) +
  theme_bw() + xlab('BW around 10,000 inhabitants') +
  ylab('Effect on spending\n(EUR / capita)') +
  theme(axis.text.x = element_text(angle = 90, vjust = 0.5, hjust=1))
p1

## Plot effect on transfers

plot_paper <- out_sens_p %>% 
  filter(period < 4) %>% 
  filter(outcome == 'ssz_total') %>% 
  mutate(period = paste0('Years included pre/post treatment: ', period)) %>% 
  filter(sample == "Full")

# Figure A.5 - Effect on per-capita transfers ----

p2 <- ggplot(plot_paper %>% filter(str_detect(period, '3')), 
             aes(bw, estimate)) +
  geom_hline(yintercept = 0, linetype = 'dotted') +
  geom_errorbar(aes(ymin = conf.low, ymax = conf.high),
                width = 0, size = 0.3, position = pd) +
  geom_errorbar(aes(ymin = conf.low90,
                    ymax = conf.high90),
                width = 0, size = 0.8,
                position = pd) +
  geom_point(shape = 21, fill = 'white', size = 2,
             position = pd) +
  theme_bw() + xlab('BW around 10,000 inhabitants') +
  ylab('Effect on transfers\n(EUR / capita)') +
  theme(axis.text.x = element_text(angle = 90, vjust = 0.5, hjust=1)) + 
  scale_x_continuous(breaks = bw_vec)
p2

## Effect on debt

plot_paper <- out_sens_p %>% 
  filter(period < 4) %>% 
  filter(outcome == 'current_debt_total') %>% 
  mutate(period = paste0('Years included\npre/post treatment: ', period)) %>% 
  filter(sample == "Full")

# Figure A.17 - Effect on per-capita debt ----

p3 <- ggplot(plot_paper, 
             aes(bw, estimate)) +
  geom_hline(yintercept = 0, linetype = 'dotted') +
  geom_errorbar(aes(ymin = conf.low, ymax = conf.high),
                width = 0, size = 0.3, position = pd) +
  geom_errorbar(aes(ymin = conf.low90,
                    ymax = conf.high90),
                width = 0, size = 0.8,
                position = pd) +
  geom_point(shape = 21, fill = 'white', size = 2,
             position = pd) +
  theme_bw() + xlab('BW around 10,000 inhabitants') +
  ylab('Effect on debt (EUR / capita)') +
  theme(axis.text.x = element_text(angle = 90, vjust = 0.5, hjust=1)) +
  facet_wrap(~period, ncol = 3) + 
  scale_x_continuous(breaks = bw_vec)
p3

## East and West Germany

plot_paper <- out_sens_p %>% 
  filter(period < 4) %>% 
  filter(outcome == 'spending_total' & bw < 2501) %>% 
  mutate(period = paste0('Years included pre/post treatment: ', period)) %>% 
  filter(sample != "Full") %>% 
  mutate(sample = factor(sample))

# Figure A.16: Effect on spending by East/West

p1 <- ggplot(plot_paper %>% 
               filter(str_detect(period, '1')),  
             aes(bw, estimate)) +
  geom_hline(yintercept = 0, linetype = 'dotted') +
  geom_errorbar(aes(ymin = conf.low, ymax = conf.high),
                width = 0, size = 0.3, position = pd) +
  geom_errorbar(aes(ymin = conf.low90,
                    ymax = conf.high90),
                width = 0, size = 0.8,
                position = pd) +
  geom_point(shape = 21, fill = 'white', size = 2,
             position = pd) +
  theme_bw() + xlab('BW around 10,000 inhabitants') +
  ylab('Effect on spending\n(EUR / capita)') +
  theme(axis.text.x = element_text(angle = 90, vjust = 0.5, hjust=1)) +
  facet_wrap(~sample)
p1 
